import os
import time
import base64
from string import Template
from playwright.async_api import async_playwright
from mcp.server.fastmcp import FastMCP
mcp = FastMCP("CoverImageService")

# 获取当前文件所在目录
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
TMP_DIR = os.path.join(CURRENT_DIR, "temp")

HTML_TEMPLATE = """
<!DOCTYPE html>
<html lang="zh-CN">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>小红书风格文本封面图</title>
    <script src="https://cdn.tailwindcss.com"></script>
    <link href="https://cdn.jsdelivr.net/npm/font-awesome@4.7.0/css/font-awesome.min.css" rel="stylesheet">
    <style type="text/tailwindcss">
        @layer utilities {
            .text-shadow {
                text-shadow: 0 2px 4px rgba(0, 0, 0, 0.1);
            }
            .glow {
                filter: drop-shadow(0 0 8px rgba(34, 197, 94, 0.5));
            }
            .highlighted {
                color: #22C55E;
                font-weight: bold;
            }
        }
    </style>
</head>
<body class="bg-gray-100 min-h-screen flex flex-col items-center justify-center p-4 gap-8">
    <!-- 封面预览（优化为直接显示，无需交互元素） -->
    <div id="cover-preview" class="relative w-full max-w-md aspect-[3/4] bg-dark rounded-2xl overflow-hidden shadow-xl">
        <div class="absolute inset-0 bg-gradient-to-b from-dark/80 to-dark/95"></div>
        
        <div class="relative h-full w-full flex flex-col items-center justify-center p-8 text-center">
            <div class="mb-4">
                <span class="text-primary text-4xl glow">🌟</span>
            </div>
            <h1 class="text-light text-[clamp(1.5rem,5vw,2.5rem)] font-bold leading-tight text-shadow mb-6">
                <div class="mb-2">
                    <span id="preview-title" class="highlightable">$highlighted_title</span>
                </div>
                <div>
                    <span id="preview-subtitle" class="highlightable">$highlighted_subtitle</span>
                </div>
            </h1>
            <p class="text-light/90 text-[clamp(1.1rem,3vw,1.5rem)] font-medium text-shadow mt-2">
                <span id="preview-description" class="highlightable">$highlighted_description</span>
            </p>
            <div class="mt-6">
                <span class="text-primary text-4xl glow">🌟</span>
            </div>
        </div>
        
        <div class="absolute bottom-4 left-0 right-0 flex justify-center">
            <div class="h-1 w-20 bg-primary/30 rounded-full"></div>
        </div>
    </div>
</body>
</html>
"""


class CoverGenerator:
    def __init__(self):
        self.output_dir = os.path.join(TMP_DIR , "generated_covers")
        self.html_temp_dir = os.path.join(TMP_DIR , "html_temp")
        print(self.output_dir)
        
        # 创建必要的目录
        for dir_path in [self.output_dir, self.html_temp_dir]:
            if not os.path.exists(dir_path):
                os.makedirs(dir_path)
        
        # HTML模板字符串（优化了截图区域）
        
    def _highlight_text(self, text, words):
        """对文本中的指定词语添加高亮标记"""
        if not words:
            return text
        
        for word in words.split(","):
            if word:
                text = text.replace(word, f"<span class=\"highlighted\">{word}</span>")
        return text
    
    def _generate_html(self, title, subtitle, description, title_highlights=None, 
                     subtitle_highlights=None, description_highlights=None, save_html=True):
        """生成包含变量替换的HTML，并可选保存为文件"""
        # 应用高亮
        highlighted_title = self._highlight_text(title, title_highlights)
        highlighted_subtitle = self._highlight_text(subtitle, subtitle_highlights)
        highlighted_description = self._highlight_text(description, description_highlights)
        
        # 替换模板变量
        template = Template(HTML_TEMPLATE)
        html_content = template.substitute(
            title=title,
            subtitle=subtitle,
            description=description,
            highlighted_title=highlighted_title,
            highlighted_subtitle=highlighted_subtitle,
            highlighted_description=highlighted_description
        )
        
        return html_content
    
    async def _generate_image(self, html_content, output_filename=None):
        """使用Playwright生成图片并返回二进制数据"""
        # 生成唯一输出文件名
        if output_filename is None:
            output_filename = f"cover_{int(time.time())}.jpg"
        
        output_path = os.path.join(self.output_dir, output_filename)
        async with async_playwright() as p:
            # 启动浏览器
            browser = await p.chromium.launch()
            page = await browser.new_page()
            
            # 设置页面内容
            await page.set_content(html_content)
            
            # 等待页面加载完成
            await page.wait_for_load_state('networkidle')

            element = page.locator('#cover-preview')
            # 获取图像二进制数据
            image_bytes = await element.screenshot()
            
            # 关闭浏览器
            await browser.close()
            
            # 保存图片到本地
            with open(output_path, 'wb') as f:
                f.write(image_bytes)
            return output_path
    
    def generate_cover(self, title, subtitle, description, 
                      title_highlights=None, subtitle_highlights=None, description_highlights=None):
        """生成封面图主函数"""
        # 生成HTML内容
        html_content = self._generate_html(
            title, subtitle, description,
            title_highlights, subtitle_highlights, description_highlights
        )
        
        # 生成图片
        image_path = self._generate_image(html_content)
        
        return image_path

# 使用示例
@mcp.tool()
async def generate_cover(title: str, subtitle: str, description: str, 
                      title_highlights: str = None, subtitle_highlights: str = None, description_highlights: str = None) -> str:
    """将文本转换为小红书风格封面图
    Args:
        title: 标题
        subtitle: 副标题
        description: 描述
        title_highlights: 标题高亮词语（逗号分隔）
        subtitle_highlights: 副标题高亮词语（逗号分隔）
        description_highlights: 描述高亮词语（逗号分隔）
    Returns:
        str: 本地图片文件路径
    """
    generator = CoverGenerator()
    html_content = generator._generate_html(
            title, subtitle, description,
            title_highlights, subtitle_highlights, description_highlights
        )
    output_path = await generator._generate_image(html_content)
    return output_path
 
if __name__ == "__main__":
    # # 生成参数
    # title = "25岁的李白"
    # subtitle = "一鸣惊人"
    # description = "青莲居士"
    # title_highlights = "李白"
    # description_highlights = "居士"
    
    # # 生成封面图
    # image_path = generate_cover(
    #     title, subtitle, description,
    #     title_highlights, None, description_highlights
    # )
    
    # print(f"封面图已生成: {image_path}")

 


    mcp.run(transport="stdio")


